import pandas as pd
import numpy as np
import os
from sklearn.model_selection import KFold
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import roc_auc_score
import warnings
warnings.filterwarnings('ignore')

print("=== 客户购买预测模型训练和预测 ===")

# 检查数据文件是否存在，如果不存在则生成模拟数据
def load_or_generate_data():
    if not os.path.exists('train_set.csv') or not os.path.exists('test_set.csv'):
        print("数据文件 train_set.csv 或 test_set.csv 未找到，生成模拟数据用于演示...")
        np.random.seed(42)
        n_samples = 1000
        
        # 生成模拟训练数据
        train = pd.DataFrame({
            'ID': range(1, n_samples + 1),
            'age': np.random.randint(18, 70, n_samples),
            'balance': np.random.randint(-5000, 50000, n_samples),
            'day': np.random.randint(1, 31, n_samples),
            'duration': np.random.randint(0, 1000, n_samples),
            'campaign': np.random.randint(1, 10, n_samples),
            'pdays': np.random.randint(-1, 100, n_samples),
            'previous': np.random.randint(0, 10, n_samples),
            'job': np.random.choice(['management', 'technician', 'admin', 'services'], n_samples),
            'marital': np.random.choice(['married', 'divorced', 'single'], n_samples),
            'y': np.random.choice([0, 1], n_samples, p=[0.8, 0.2])
        })
        
        # 生成模拟测试数据
        test = pd.DataFrame({
            'ID': range(n_samples + 1, n_samples + 501),
            'age': np.random.randint(18, 70, 500),
            'balance': np.random.randint(-5000, 50000, 500),
            'day': np.random.randint(1, 31, 500),
            'duration': np.random.randint(0, 1000, 500),
            'campaign': np.random.randint(1, 10, 500),
            'pdays': np.random.randint(-1, 100, 500),
            'previous': np.random.randint(0, 10, 500),
            'job': np.random.choice(['management', 'technician', 'admin', 'services'], 500),
            'marital': np.random.choice(['married', 'divorced', 'single'], 500)
        })
        
        train.to_csv('train_set.csv', index=False)
        test.to_csv('test_set.csv', index=False)
        print(f"已生成模拟数据: 训练集 {len(train)} 条, 测试集 {len(test)} 条")
        return train, test
    else:
        train = pd.read_csv('train_set.csv')
        test = pd.read_csv('test_set.csv')
        print("成功加载数据文件")
        return train, test

# 特征工程
def feature_engineering(train, test):
    print("\n=== 特征工程 ===")
    
    # 处理异常值
    def outlier_processing(series):
        series_clean = series.copy()
        q1 = series_clean.quantile(0.25)
        q3 = series_clean.quantile(0.75)
        iqr = q3 - q1
        lower_bound = q1 - 1.5 * iqr
        upper_bound = q3 + 1.5 * iqr
        return series_clean.clip(lower=lower_bound, upper=upper_bound)
    
    # 对数值型特征处理异常值
    num_features = ['age', 'balance', 'day', 'duration', 'campaign', 'pdays', 'previous']
    for col in num_features:
        if col in train.columns:
            train[col] = outlier_processing(train[col])
        if col in test.columns:
            test[col] = outlier_processing(test[col])
    
    # 对分类变量进行one-hot编码
    str_features = ['job', 'marital']
    for col in str_features:
        if col in train.columns and col in test.columns:
            all_categories = set(train[col].unique()) | set(test[col].unique())
            for category in all_categories:
                train[f'{col}_{category}'] = (train[col] == category).astype(int)
                test[f'{col}_{category}'] = (test[col] == category).astype(int)
    
    # 准备特征矩阵
    features_to_drop = ['ID', 'y'] + str_features
    X = train.drop([col for col in features_to_drop if col in train.columns], axis=1)
    y = train['y']
    
    test_features = test.drop([col for col in ['ID'] + str_features if col in test.columns], axis=1)
    
    # 标准化数值特征
    numeric_cols = [col for col in X.columns if col not in [f'{s}_{c}' for s in str_features for c in ['management', 'technician', 'admin', 'services', 'married', 'divorced', 'single']]]
    if numeric_cols:
        scaler = StandardScaler()
        X[numeric_cols] = scaler.fit_transform(X[numeric_cols])
        test_features[numeric_cols] = scaler.transform(test_features[numeric_cols])
    
    print(f"特征工程完成，训练特征形状: {X.shape}, 测试特征形状: {test_features.shape}")
    return X, y, test_features

# 模型训练和预测
def train_and_predict(X, y, test_features):
    print("\n=== 模型训练和预测 ===")
    
    # 使用随机森林进行5折交叉验证
    n_fold = 5
    folds = KFold(n_splits=n_fold, shuffle=True, random_state=2022)
    oof_predictions = np.zeros(len(X))
    test_predictions = np.zeros(len(test_features))
    
    for fold_n, (train_index, valid_index) in enumerate(folds.split(X)):
        print(f"训练第 {fold_n + 1}/{n_fold} 折...")
        
        X_train, X_valid = X.iloc[train_index], X.iloc[valid_index]
        y_train, y_valid = y[train_index], y[valid_index]
        
        # 使用随机森林回归器（输出概率值）
        model = RandomForestRegressor(
            n_estimators=100,  # 减少树的数量以加快训练速度
            max_depth=10,
            random_state=2022,
            n_jobs=-1
        )
        
        model.fit(X_train, y_train)
        
        # 验证集预测
        y_pred_valid = model.predict(X_valid)
        oof_predictions[valid_index] = y_pred_valid
        
        # 测试集预测
        test_pred = model.predict(test_features)
        test_predictions += test_pred / n_fold
    
    # 计算模型性能
    auc_score = roc_auc_score(y, oof_predictions)
    print(f"模型性能 - AUC Score: {auc_score:.6f}")
    
    return test_predictions, auc_score

# 生成提交文件
def generate_submission(test, predictions, auc_score):
    print("\n=== 生成预测结果 ===")
    
    submission = pd.DataFrame({
        'ID': test['ID'],
        'pred': predictions
    })
    
    submission_file = f'customer_purchase_predictions_auc_{auc_score:.6f}.csv'
    submission.to_csv(submission_file, index=False)
    print(f"预测结果已保存到: {submission_file}")
    print(f"预测样本数量: {len(submission)}")
    print("前5个预测结果:")
    print(submission.head())
    
    return submission_file

# 主函数
def main():
    # 加载数据
    train, test = load_or_generate_data()
    
    # 特征工程
    X, y, test_features = feature_engineering(train, test)
    
    # 模型训练和预测
    predictions, auc_score = train_and_predict(X, y, test_features)
    
    # 生成提交文件
    submission_file = generate_submission(test, predictions, auc_score)
    
    print(f"\n=== 任务完成 ===")
    print(f"最终AUC分数: {auc_score:.6f}")
    print(f"预测文件: {submission_file}")

if __name__ == "__main__":
    main()